"""
python p2_solution_aco_ant.py
"""

import json
import os
import networkx as nx
import random

# 路径设置
warehouse_path = '../fujian/fujian3/origin_data/warehouse.json'
inventory_path = '../fujian/fujian3/data_from_p1/all_average_inventory.json'
sales_path = '../fujian/fujian3/data_from_p1/all_average_sales.json'
output_path = '../fujian/fujian3/ant_algo/'

# 数据加载
with open(warehouse_path, 'r') as f:
    warehouses = json.load(f)
with open(inventory_path, 'r') as f:
    inventory_data = json.load(f)
with open(sales_path, 'r') as f:
    sales_data = json.load(f)

# 数据处理
warehouse_dict = {w['warehouse_id']: w for w in warehouses}
inventory_dict = {i['category_id']: i['average_inventory'] for i in inventory_data}
sales_dict = {s['category_id']: s['average_sales'] for s in sales_data}

# 创建网络图
graph = nx.Graph()

# 添加仓库节点
for warehouse_id, warehouse_info in warehouse_dict.items():
    graph.add_node(warehouse_id)

# ACO参数设置
class Ant:
    def __init__(self, warehouse_dict, inventory_dict, sales_dict):
        self.warehouse_dict = warehouse_dict
        self.inventory_dict = inventory_dict
        self.sales_dict = sales_dict
        self.assignment = {}

    # def build_solution(self):
    #     for category in self.inventory_dict.keys():
    #         warehouse = random.choice(list(self.warehouse_dict.keys()))
    #         self.assignment[category] = warehouse
    #     return self.assignment
    
    def build_solution(self):
        for category in self.inventory_dict.keys():
            # 优先选择能满足当前类别库存的仓库
            available_warehouses = [
                warehouse_id for warehouse_id in self.warehouse_dict.keys()
                if self.warehouse_dict[warehouse_id]['max_inventory'] >= self.inventory_dict[category]
            ]
            if available_warehouses:
                warehouse = random.choice(available_warehouses)
                self.assignment[category] = warehouse
            else:
                # 如果没有可用仓库，随机选择
                warehouse = random.choice(list(self.warehouse_dict.keys()))
                self.assignment[category] = warehouse
        return self.assignment


    def evaluate_solution(self):
        total_cost = 0
        warehouse_utilization = {w: {'inventory': 0, 'sales': 0} for w in self.warehouse_dict.keys()}

        for category, warehouse in self.assignment.items():
            inventory = self.inventory_dict[category]
            sales = self.sales_dict[category]
            warehouse_utilization[warehouse]['inventory'] += inventory
            warehouse_utilization[warehouse]['sales'] += sales

        # 检查约束条件
        for warehouse_id, usage in warehouse_utilization.items():
            if (usage['inventory'] >= self.warehouse_dict[warehouse_id]['max_inventory'] or
                usage['sales'] >= self.warehouse_dict[warehouse_id]['max_sales']):
                return float('inf')  # 违反约束条件

        # 满足条件时计算总成本
        total_cost = sum(self.warehouse_dict[w]['daily_cost'] for w in warehouse_utilization if warehouse_utilization[w]['inventory'] > 0)
        return total_cost

# ACO主逻辑
def solve_warehouse_assignment(num_ants=50, iterations=100):
    best_solution = None
    best_cost = float('inf')

    for _ in range(iterations):
        for _ in range(num_ants):
            ant = Ant(warehouse_dict, inventory_dict, sales_dict)
            solution = ant.build_solution()
            cost = ant.evaluate_solution()

            if cost < best_cost and cost != float('inf'):
                best_cost = cost
                best_solution = solution

    return best_solution, best_cost

# 生成解决方案并保存结果
solution, total_cost = solve_warehouse_assignment()

# 检查解决方案是否有效
if solution is None:
    print("没有找到有效的解决方案。")
else:
    warehouse_utilization = {w: {'inventory': 0, 'sales': 0} for w in warehouse_dict.keys()}

    for category, warehouse in solution.items():
        inventory = inventory_dict[category]
        sales = sales_dict[category]
        warehouse_utilization[warehouse]['inventory'] += inventory
        warehouse_utilization[warehouse]['sales'] += sales

    # 计算利用率
    utilization_rates_inventory = []
    utilization_rates_sales = []

    for warehouse_id, usage in warehouse_utilization.items():
        utilization_rates_inventory.append({
            'warehouse_id': warehouse_id,
            'utilization_rate_of_inventory': usage['inventory'] / warehouse_dict[warehouse_id]['max_inventory']
        })
        utilization_rates_sales.append({
            'warehouse_id': warehouse_id,
            'utilization_rate_of_sales': usage['sales'] / warehouse_dict[warehouse_id]['max_sales']
        })

    # 输出文件夹
    os.makedirs(output_path, exist_ok=True)

    # 保存总成本
    with open(os.path.join(output_path, 'all_cost.txt'), 'w') as f:
        f.write(str(total_cost))

    # 保存库存利用率
    with open(os.path.join(output_path, 'all_utilization_rate_inventory.json'), 'w') as f:
        json.dump(utilization_rates_inventory, f)

    # 保存出货量利用率
    with open(os.path.join(output_path, 'all_utilization_rate_sales.json'), 'w') as f:
        json.dump(utilization_rates_sales, f)

    # 保存分配方案
    with open(os.path.join(output_path, 'assignment.json'), 'w') as f:
        json.dump([{'category_id': k, 'warehouse_id': v} for k, v in solution.items()], f)
